#ifndef KPEXTVKCONTEXT_H
#define KPEXTVKCONTEXT_H

#include "kpextCommonDef.h"
#include "kpextCommandBuffer.h"
#include "kpextMemoryResource.h"
#include "kpextDescriptorMgr.h"
#include "kpextComputePipeline.h"

#define KPEXT_CTX_NOK_S "Vulkan context is not initialize success!"

#define KPEXT_CTX_CHECK(T) \
    do{ \
    if(!getIsVkContextOk()){ \
        SPDLOG_ERROR(KPEXT_CTX_NOK_S); \
        return std::shared_ptr<T>(nullptr); \
    } \
    }while(0) \


#define KPEXT_DECLARE_SPIRV(name, spvdata_uint) \
    std::vector<uint32_t> name(spvdata_uint, \
    spvdata_uint + sizeof(spvdata_uint) / sizeof(spvdata_uint[0]) \
    ) \


/*****************/

/**
 * @brief 根据给定的workgroupsize工作组尺寸，和二维输入行、列， 计算工作组
 * @param cols
 * @param rows
 * @param workgroupsize
 * @return
 */
std::array<uint32_t, 3> calc2DWorkGroup(uint32_t cols, uint32_t rows,
                                        const std::array<uint32_t, 3> workgroupsize);


/**
 * @brief vulkan 上下文对象, 用于创建vulkan instance,physicaldevice,device,queque
 */
class kpextVkContext
{
public:
    explicit kpextVkContext();
    virtual ~kpextVkContext();

    /**
     * @brief 销毁vulkan实例对象， 本函数是非线程安全的， 请确保所有instance相关的vulkan资源已经销毁再调用
     */
    static bool destroyInstance();


    /**
     * @brief 返回vulkan上下文是否正常完整
     * @return  true 表示  vulkan上下文正常完整
     */
    virtual bool getIsVkContextOk();


    /**
     * @brief 获取默认的glsl计算着色器WorkGroupSize
     * @return
     */
    virtual std::array<uint32_t, 3> getDefWorkGroupSize();

    /**
     * @brief 获取当前vulkan物理设备计算着色器工作组数量最大值
     * @return
     */
    virtual std::array<uint32_t, 3> getMaxWorkGroupCount();

    /**
     * @brief 获取当前vulkan逻辑设备建立的队列数量
     * @return  队列数量
     */
    virtual uint32_t getComputeQueuesCount();

    /**
     * @brief 创建kpext::CommandBuffer对象， 只是vulkan命令池和命令缓冲的简单包装类， 本函数只能创建主要缓冲区
     * @param queueIndex 队列索引, 用于查找创建vukan设备一并创建的命令队列
     * @param poolflag 与命令缓冲一起创建的命令池属性标记
     * @param totalTimestamps  参考kp::Manager::sequence函数
     * @return 返回kpext::CommandBuffer对象共享指针
     */
    std::shared_ptr<kpext::CommandBuffer> cmdbuf(uint32_t queueIndex = 0,
                                                 vk::CommandPoolCreateFlags poolflag = vk::CommandPoolCreateFlagBits::eResetCommandBuffer,
                                                 uint32_t totalTimestamps = 0);

    /**
     * @brief 创建kpext::RawBuffer对象， 只是vulkan Buffer简单包装， (未创建vulkan Buffer)
     * @return 返回kpext::RawBuffer对象共享指针
     */
    std::shared_ptr<kpext::RawBuffer> buffer();

    /**
     * @brief 创建kpext::RawDeviceMemory对象， 只是vulkan DeviceMemory简单包装， (未创建vulkan DeviceMemory)
     * @return 返回kpext::RawDeviceMemory对象共享指针
     */
    std::shared_ptr<kpext::RawDeviceMemory> deviceMem();

    /**
     * @brief 创建kpext::BaseBuffer对象，是kpext::RawBuffer和kpext::kpext::RawDeviceMemory简单集合类，创建后应调用rebuild创建显存
     * @param size 缓冲区字节大小
     * @return  kpext::BaseBuffer对象共享指针
     */
    std::shared_ptr<kpext::BaseBuffer> baseBuffer(size_t size);


    /**
     * @brief 创建kpext::StagingBuffer 暂存缓冲对象
     * @param size  缓冲区的节字大小
     * @param coherent  缓冲区是否与主机端保持一致性（一般用于主机传数据到vulkan设备）
     * @param cached   缓冲区是否在主机端cache缓冲（一般用于vulkan设备传数据到主机, 需要配合flushMapped invalidateMapped使用）
     * @return kpext::StagingBuffer 暂存缓冲对象共享指针
     */
    std::shared_ptr<kpext::StagingBuffer> stagingBuf(
            size_t size = 4,
            bool coherent = true,
            bool cached = false
            );


    /**
     * @brief 创建kpext::DeviceBuffer 仅vulkan逻辑设备可见的缓冲对象
     * @param size  缓冲区的节字大小
     * @return kpext::DeviceBuffer 设备缓冲对象共享指针
     */
    std::shared_ptr<kpext::DeviceBuffer> deviceBuf(size_t size = 4);

    /**
     * @brief 创建kpext::UniformBuffer 仅vulkan逻辑设备可见的只读缓冲对象, 着色器里的代码访问同一内存
     * @param size 缓冲区的节字大小
     * @return kpext::UniformBuffer 统一缓冲对象共享指针
     */
    std::shared_ptr<kpext::UniformBuffer> uniformBuf(size_t size = 4);

    /**
     * @brief 检查统一纹素缓冲区的 元素索引大小 是否能够使用(未超出最大索引)
     * @param size 元素索引大小
     * @return 能使用返回true, 否则返回false
     */
    bool checkUniformTexelBufElemSize(uint32_t size);

    /**
     * @brief 创建kpext::UniformTexelBuffer 仅vulkan逻辑设备可见的只读纹素缓冲对象, 着色器里的代码访问同一内存
     * @param size 缓冲区的节字大小
     * @param  format 纹素格式
     * @return kpext::UniformBuffer 统一缓冲对象共享指针
     */
    std::shared_ptr<kpext::UniformTexelBuffer> uniformTexelBuf(size_t size = 4,
                                                               vk::Format format = vk::Format::eR32Sfloat);


    /**
     * @brief 创建vulkan 描述符池包装类对象, {存储缓冲,uniform纹素缓冲,uniform缓冲,纹素缓冲,存储图像}描述符数量向量
     * @param poolCap 描述符池最多可建立的描述符容量数组，{存储缓冲,uniform纹素缓冲,uniform缓冲,纹素缓冲,存储图像}描述符数量向量
     * @param maxSetsCount  描述符池最多可建立的描述符集
     * @param flags 描述符池创建特性，可指定是否能回收资源、是否能绑定后更新描述符集
     * @return 成功返回vulkan 描述符池包装类对象共享指针， 否则返回空指针
     */
    std::shared_ptr<kpext::DescriptorPool>
    descriptorPool( const std::vector<uint32_t> & poolCap = {1},
                    uint32_t maxSetsCount = 1,
                    vk::DescriptorPoolCreateFlags flags = vk::DescriptorPoolCreateFlags());


    /**
     * @brief 创建一个kpext::ComputePipeline类对象，对象内部未创建vulkan pipeline
     * @return kpext::ComputePipeline类对象共享指针
     */
    std::shared_ptr<kpext::ComputePipeline> computePipeline();


#ifdef KPEXT_USE_KOMPUTE
    /**
     * @brief 创建kp::Sequence对象，并返回其共享指针，不受本上下文类管理(不自动释放)
     * @param queueIndex  队列索引, 用于查找创建vukan设备一并创建的命令队列
     * @param totalTimestamps  参考kp::Manager::sequence函数
     * @return 返回kp::Sequence对象共享指针
     */
    std::shared_ptr<kp::Sequence> sequence(uint32_t queueIndex = 0,
                                           uint32_t totalTimestamps = 0);



    /**
     * ****************************************************
     * 由于目前与kp::Tensor相关的实现其内部都创建vkBuffer对象的，
     * 并且使用vk::SharingMode::eExclusive独占模式，
     * 这种模式由首次调用它的queueFamily独占，如果要在其它queueFamily必须通过VkBufferMemoryBarrier
     * 进行权属转移，而kompute库并没有考虑这样的问题，所以它创建的Tensor应尽量只在一个queueFamily中使用，
     * kp::OpMemoryBarrier 类也没有进行权属转移相关设定
     * ****************************************************
    */


    /**
     * @brief 与参考kp::Manager::tensorT<T>功能相同，返回的共享指针不受本上下文类管理(不自动释放)
     * @param data 用于初始化的数据向量
     * @param tensorType  一般用kp::Tensor::TensorTypes::eDevice就可以了
     * @return 返回kp::TensorT<T>对象的共享指针
     */
    template<typename T>
    std::shared_ptr<kp::TensorT<T>> tensorT(
            const std::vector<T>& data,
            kp::Tensor::TensorTypes tensorType = kp::Tensor::TensorTypes::eDevice){
        KPEXT_CTX_CHECK(kp::TensorT<T>);

        SPDLOG_DEBUG("kp::Tensor<T> creation triggered");

        std::shared_ptr<kp::TensorT<T>> tensor{ new kp::TensorT<T>(
                        this->m_phyDev, this->m_dev, data, tensorType) };

        return tensor;
    }


    /**
     * @brief tensorT的浮点数版本, 返回的共享指针不受本上下文类管理(不自动释放)
     * @param data 用于初始化的数据向量
     * @param tensorType 一般用kp::Tensor::TensorTypes::eDevice就可以了
     * @return 返回kp::TensorT<float>对象的共享指针
     */
    virtual std::shared_ptr<kp::TensorT<float>> tensor(
            const std::vector<float>& data,
            kp::Tensor::TensorTypes tensorType = kp::Tensor::TensorTypes::eDevice);


    /**
     * @brief 与kp::Manager::tensor函数功能相同， 返回的共享指针不受本上下文类管理(不自动释放)
     * @param data  用于初始化的数据指针，不能为空指针且必须包含elementTotalCount * elementMemorySize节字数据
     * @param elementTotalCount  data包含多少个元素
     * @param elementMemorySize  data每个元素占用的字节数
     * @param dataType  data的元素类型， 目前可用的为eBool,eInt,eUnsignedInt,eFloat,eDouble
     * @param tensorType  一般用kp::Tensor::TensorTypes::eDevice就可以了
     * @return  返回kp::Tensor对象的共享指针
     */
    virtual std::shared_ptr<kp::Tensor> tensor(
            void* data,
            uint32_t elementTotalCount,
            uint32_t elementMemorySize,
            const kp::Tensor::TensorDataTypes& dataType,
            kp::Tensor::TensorTypes tensorType = kp::Tensor::TensorTypes::eDevice);


    /**
     * @brief 创建二维图像张量, 类型为单通道32位浮点
     * @param cols 二维图像的列数
     * @param rows 二维图像的行数
     * @param data 初始化数据，如果为null，全部初始化为零
     * @return 返回张量对象指针
     */
    virtual std::shared_ptr<kp::Tensor> imgF32C1(int cols, int rows,
                                                 void * data = nullptr);

    /**
     * @brief 创建二维图像张量, 类型为单通道16位无符号数(返回的是uint32的tensor)
     * @param cols  二维图像的列数
     * @param rows  二维图像的行数
     * @param data  初始化数据，如果为null，全部初始化为零. data的个数应保证 cols * rows 向上取整的偶数个
     * @param keepeven 是否创建tensor元素个数始终向上取整为偶数个
     * @return  返回张量对象指针
     */
    virtual std::shared_ptr<kp::Tensor> imgU16C1(int cols, int rows,
                                                 void * data = nullptr,
                                                 bool keepeven = true
                                                 );

    /**
     * @brief 创建一个kp::Algorithm算法对象, 特化常量和推送常量类型模板化，返回共享指针，返回的共享指针不受本上下文类管理(不自动释放)
     * @param tensors  输入到计算着色器的张量数组
     * @param spirv  由着色器代码编译出的标准中间表示代码vector<uint32>
     * @param workgroup  dispatch到计算着色器的workgroup, 参考glsl计算着色器
     * @param specializationConstants  输入到计算着色器的特化常量
     * @param pushConstants  输入到计算着色器的推送常量
     * @return  kp::Algorithm算法对象共享对象
     */
    template<typename S = float, typename P = float>
    std::shared_ptr<kp::Algorithm> algorithm(
            const std::vector<std::shared_ptr<kp::Tensor>>& tensors,
            const std::vector<uint32_t>& spirv,
            const kp::Workgroup& workgroup,
            const std::vector<S>& specializationConstants,
            const std::vector<P>& pushConstants)
    {
        KPEXT_CTX_CHECK(kp::Algorithm);

        SPDLOG_DEBUG("kp::Algorithm<S, P> creation triggered");

        std::shared_ptr<kp::Algorithm> algorithm{ new kp::Algorithm(
                        this->m_dev,
                        tensors,
                        spirv,
                        workgroup,
                        specializationConstants,
                        pushConstants) };

        return algorithm;
    }

    /**
     * @brief 创建一个kp::Algorithm算法对象, 特化常量和推送常量类型为单精度浮点，返回共享指针，返回的共享指针不受本上下文类管理(不自动释放)
     * @param tensors  输入到计算着色器的张量数组
     * @param spirv  由着色器代码编译出的标准中间表示代码vector<uint32>
     * @param workgroup  dispatch到计算着色器的workgroup, 参考glsl计算着色器
     * @param specializationConstants  输入到计算着色器的特化常量
     * @param pushConstants  输入到计算着色器的推送常量
     * @return  kp::Algorithm算法对象共享对象
     */
    virtual std::shared_ptr<kp::Algorithm> algorithm(
            const std::vector<std::shared_ptr<kp::Tensor>>& tensors = {},
            const std::vector<uint32_t>& spirv = {},
            const kp::Workgroup& workgroup = {},
            const std::vector<float>& specializationConstants = {},
            const std::vector<float>& pushConstants = {});

    /**
     * @brief 创建一个kp::Algorithm算法对象, 特化常量和推送常量类型为32位整数，返回共享指针，返回的共享指针不受本上下文类管理(不自动释放)
     * @param tensors  输入到计算着色器的张量数组
     * @param spirv  由着色器代码编译出的标准中间表示代码vector<uint32>
     * @param workgroup  dispatch到计算着色器的workgroup, 参考glsl计算着色器
     * @param specializationConstants  输入到计算着色器的特化常量
     * @param pushConstants  输入到计算着色器的推送常量
     * @return  kp::Algorithm算法对象共享对象
     */
    virtual std::shared_ptr<kp::Algorithm> algorithm(
            const std::vector<std::shared_ptr<kp::Tensor>>& tensors = {},
            const std::vector<uint32_t>& spirv = {},
            const kp::Workgroup& workgroup = {},
            const std::vector<int32_t>& specializationConstants = {},
            const std::vector<int32_t>& pushConstants = {});

    /**
     * @brief 创建一个kp::Algorithm算法对象, 特化常量和推送常量类型为32位无符号整数，返回共享指针，返回的共享指针不受本上下文类管理(不自动释放)
     * @param tensors  输入到计算着色器的张量数组
     * @param spirv  由着色器代码编译出的标准中间表示代码vector<uint32>
     * @param workgroup  dispatch到计算着色器的workgroup, 参考glsl计算着色器
     * @param specializationConstants  输入到计算着色器的特化常量
     * @param pushConstants  输入到计算着色器的推送常量
     * @return  kp::Algorithm算法对象共享对象
     */
    virtual std::shared_ptr<kp::Algorithm> algorithm(
            const std::vector<std::shared_ptr<kp::Tensor>>& tensors = {},
            const std::vector<uint32_t>& spirv = {},
            const kp::Workgroup& workgroup = {},
            const std::vector<uint32_t>& specializationConstants = {},
            const std::vector<uint32_t>& pushConstants = {});

#endif

#ifndef KOMPUTE_DISABLE_VK_DEBUG_LAYERS
    /**
     * @brief 是否使能vulkan 验证层标记
     */
    static bool EnableDebugLayers;
    /**
     * @brief 是否使能vulkan 验证层标记已经由程序进行了设置
     */
    static bool EnableDebugLayersByProgram;
#endif

    /**
     * @brief 禁用 json文件配置
     */
    static bool NO_USE_JSON_CONFIG;

    /**
     * @brief 期望创建使用的vulkan物理设备索引
     */
    static int VK_PHY_DEV_IDX_TO_CREATE;

    /**
     * @brief 期望创建使用的vulkan物理设备pipelineuuid, VK_PHY_DEV_IDX_TO_CREATE为-1时才有效
     */
    static std::string VK_PHY_DEV_PUUID_TO_CREATE;

    /**
     * @brief 期望创建使用的队列家族
     */
    static std::vector<unsigned int> VK_QUEUE_FAMILY_INDICES;

    /**
     * @brief 期望使用的逻辑设备扩展功能
     */
    static std::vector<std::string> VK_DESIRED_DEV_EXTENSIONS;

    /**
     * @brief 默认的计算着色器尺寸
     */
    static std::vector<unsigned int> VK_DEF_COMPUTESHADER_WORKGROUP_SIZE;

    /**
     * @brief 各阶段图像处理算法需要的线程数
     */
    static std::vector<unsigned int> StagesThreadsCount;

protected:

    /**
     * @brief json配置信息
     */
    nlohmann::json * m_ctxCfg = new nlohmann::json{};

    /**
     * @brief vulkan instance
     */
    static std::shared_ptr<vk::Instance> m_instance;

#ifndef KOMPUTE_DISABLE_VK_DEBUG_LAYERS
    static vk::DebugUtilsMessengerEXT m_debugMessenger;
    static vk::DispatchLoaderDynamic m_dynamicDispatcher;
#endif

    std::string m_cfgFilePath{"kpextVkContextCfg.json"};

    /**
     * @brief vulkan 物理设备
     */
    std::shared_ptr<vk::PhysicalDevice> m_phyDev = nullptr;

    /**
     * @brief vulkan 逻辑设备
     */
    std::shared_ptr<vk::Device> m_dev = nullptr;

    /**
     * @brief 支持计算着色器的队列家族索引数组
     */
    std::vector<unsigned int> m_ComputeQueueFamilyIndices{};

    /**
     * @brief 与vulkan 逻辑设备一并创建的命令队列对象
     */
    std::vector<std::shared_ptr<vk::Queue>> m_ComputeQueues{};

    /**
     * @brief 与vulkan 逻辑设备一并创建的命令队列对象的互斥锁
     */
    std::vector<std::shared_ptr<std::mutex>> m_mutexComputeQueues{};

    /**
     * @brief vulkan 物理设备支持的显存特性
     */
    vk::PhysicalDeviceMemoryProperties m_supportMemProps = vk::PhysicalDeviceMemoryProperties();


    std::shared_ptr<vk::PipelineCache> m_GlobaPipelinecache = nullptr;

    /**
     * @brief 读取json配置信息
     */
    virtual void readJsonCfg();

    /**
     * @brief 从json配置对象中读取要创建的vulkan物理对象pipelineuuid信息
     * @return  返回pipelineuuid信息(转化为HEX之后的UUID)
     */
    virtual std::string getPipelineUUIDFromJsonCfg();

    /**
     * @brief 从json配置对象中读取要创建的vulkan设备的队列家族索引数组
     * @return 要创建的设备的队列家族索引数组
     */
    virtual std::vector<unsigned int> getQueueFamilyIndicesFromJsonCfg();

    /**
     * @brief 向json对象中更新vulkan配置对象信息
     * @param jo 目标json对象
     * @param phydevs 要查询的vulkan物理设备对象数组
     */
    virtual void updateJsonPhyDevInfo(nlohmann::json & jo,
                                      std::vector<vk::PhysicalDevice> & phydevs);

    /**
     * @brief 保存JSON对象到文件中
     * @param jo  要保存的json对象
     * @param filename  要保存的文件名
     */
    virtual void saveJsonToFile(nlohmann::json & jo, std::string filename);

    /**
     * @brief 创建vulkan instance对象
     */
    virtual bool createInstance();

    /**
     * @brief 创建 vulkan 物理和逻辑设备对象、命令队列
     * @param familyQueueIndices  要创建的命令队列所属的队列家族
     * @param physicalDeviceIndex  要创建的设备索引
     * @param desiredExtensions  要创建的逻辑设备需要打开的扩展功能
     */
    virtual void createDevice(const std::vector<unsigned int>& familyQueueIndices = {},
                              unsigned int physicalDeviceIndex = 0,
                              const std::vector<std::string>& desiredExtensions = {});



};

/**
 * @brief 库内部使用的一个全局vulkan上下文对象指针
 */
extern kpextVkContext * GvkContextPtr;

#endif // KPEXTVKCONTEXT_H
